import os
import cv2
from keras.models import load_model
import numpy as np


def get_inputs(src=[]):
  pre_x = []
  for s in src:
    input = cv2.imread(s)
    input = cv2.resize(input, (150, 150))
    input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB)
    pre_x.append(input) # input一张图片
  pre_x = np.array(pre_x) / 255.0
  return pre_x


camera = cv2.VideoCapture(0)
model = load_model('model\\test2.h5')
name_list = ['限速40', '限速60', '限速80', '禁止左转', '禁止通行', '禁鸣', '靠右行驶', '机动车道', '注意行人', '禁止停车', '禁止驶入']
while True:
    ret, frame = camera.read()
    cv2.imshow('frame', frame)
    image = cv2.resize(frame, (150, 150))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    p_image = np.expand_dims(image/255.0, axis=0)
    x = model.predict(p_image)
    type_name = name_list[np.argmax(x)]
    pro = 100 * np.around(np.max(x), 4)
    print("这是 {} 标志，预测概率：{} %。".format(type_name, pro))

    # print(image.shape)
    if cv2.waitKey(1) == ord('q'):
        break

camera.release()
cv2.destroyAllWindows()




# images = []
# test_dir = 'dataset\\traffic_sign\\test'
# for fn in os.listdir(test_dir):
#     path = os.path.join(test_dir, fn)
#     print(path)
#     images.append(path)
#
# pre_x = get_inputs(images)
# pre_y = model.predict_classes(pre_x)
# print(pre_y)

